from mgamdata.dataset.base import mgam_SemiSup_3D_Mha
from mgamdata.dataset.RenJi_Sarcopenia.mm_dataset import Sarcopenia_Mha
from mgamdata.dataset.Totalsegmentator.mm_dataset import TotalsegmentatorSeg3DDataset
from mgamdata.dataset.FLARE_2023.mm_dataset import FLARE_2023_Semi_Mha
from mgamdata.dataset.RenJi_Sarcopenia.mm_dataset import Sarcopenia_Mha
from mgamdata.dataset.FLARE_2022.mm_dataset import FLARE_2022_Semi_Mha
from mgamdata.dataset.ImageTBAD.mm_dataset import TBAD_Mha
from mgamdata.dataset.LUNA16.mm_dataset import LUNA16_Mha
from mgamdata.dataset.AbdomenCT_1K.mm_dataset import AbdomenCT_1K_Semi_Mha
from mgamdata.dataset.KiTS23.mm_dataset import KiTS23_Mha

mgam_dataset_root = "/mnt/wsl/PHYSICALDRIVE4/mgam_datasets/"
Sar_mha_data_root           = mgam_dataset_root + "Sarcopenia_Data/spacing2_mha/"
Tsd_mha_data_root           = mgam_dataset_root + "Totalsegmentator_Data/spacing2_mha/"
FLARE2021_mha_data_root     = mgam_dataset_root + "FLARE_2021/spacing2_mha/"
FLARE2022_mha_data_root     = mgam_dataset_root + "FLARE_2022/spacing2_mha/"
FLARE2023_mha_data_root     = mgam_dataset_root + "FLARE_2023/spacing2_mha/"
ImageTBAD_mha_data_root     = mgam_dataset_root + "ImageTBAD/spacing2_mha/"
KiTS23_mha_data_root        = mgam_dataset_root + "KiTS23/spacing2_mha/"
LUNA16_mha_data_root        = mgam_dataset_root + "LUNA16/spacing2_mha/"
AbdomenCT1K_mha_data_root   = mgam_dataset_root + "AbdomenCT_1K/spacing2_mha/"
GastricCancer_mha_data_root = mgam_dataset_root + "GastricCancerSegmentation_2023/spacing2_mha/"
Pediatric_mha_data_root     = mgam_dataset_root + "Pediatric_CT_SEG/spacing2_mha/"
StageII_mha_data_root       = mgam_dataset_root + "StageII_Colorectal_CT/spacing2_mha/"
CPTAC_SAR_mha_data_root     = mgam_dataset_root + "CPTAC_SAR/spacing2_mha/"
Colonography_mha_data_root  = mgam_dataset_root + "CT_Colonography/spacing2_mha/"
CT_ORG_mha_data_root        = mgam_dataset_root + "CT_ORG/spacing2_mha/"

TCGA_mha_data_roots = dict(
    TCGA_STAD = mgam_dataset_root + "TCGA-STAD/spacing2_mha/",
    TCGA_UCEC = mgam_dataset_root + "TCGA-UCEC/spacing2_mha/",
    TCGA_KIRC = mgam_dataset_root + "TCGA-KIRC/spacing2_mha/",
    TCGA_BLCA = mgam_dataset_root + "TCGA-BLCA/spacing2_mha/",
    TCGA_OV   = mgam_dataset_root + "TCGA-OV/spacing2_mha/",
    TCGA_LIHC = mgam_dataset_root + "TCGA-LIHC/spacing2_mha/",
)

Totalseg_Dataset = dict(
    type=TotalsegmentatorSeg3DDataset,
    data_root=Tsd_mha_data_root,
    data_root_mha=Tsd_mha_data_root,
)
FLARE2021_Dataset = dict(
    type=FLARE_2022_Semi_Mha,
    dataset_name="FLARE_2021", 
    data_root=FLARE2021_mha_data_root,
    data_root_mha=FLARE2021_mha_data_root,
)
FLARE2022_Dataset = dict(
    type=FLARE_2022_Semi_Mha,
    data_root=FLARE2022_mha_data_root,
    data_root_mha=FLARE2022_mha_data_root,
)
FLARE2023_Dataset = dict(
    type=FLARE_2023_Semi_Mha,
    data_root=FLARE2023_mha_data_root,
    data_root_mha=FLARE2023_mha_data_root,
)
Sarcopenia_Dataset = dict(
    type=Sarcopenia_Mha, 
    data_root=Sar_mha_data_root, 
    data_root_mha=Sar_mha_data_root, 
)
ImageTBAD_Dataset = dict(
    type=TBAD_Mha,
    data_root=ImageTBAD_mha_data_root,
    data_root_mha=ImageTBAD_mha_data_root,
)
KiTS23_Dataset = dict(
    type=KiTS23_Mha,
    data_root=KiTS23_mha_data_root,
    data_root_mha=KiTS23_mha_data_root,
)
LUNA16_Dataset = dict(
    type=LUNA16_Mha,
    data_root=LUNA16_mha_data_root,
    data_root_mha=LUNA16_mha_data_root,
)
AbdomenCT1K_Dataset = dict(
    type=AbdomenCT_1K_Semi_Mha,
    data_root=AbdomenCT1K_mha_data_root,
    data_root_mha=AbdomenCT1K_mha_data_root,
)
GastricCancer_Dataset = dict(
    type=mgam_SemiSup_3D_Mha,
    dataset_name="GastricCancer",
    data_root=KiTS23_mha_data_root,
    data_root_mha=KiTS23_mha_data_root,
)
Pediatric_Dataset = dict(
    type=mgam_SemiSup_3D_Mha,
    dataset_name="Pediatric",
    data_root=Pediatric_mha_data_root,
    data_root_mha=Pediatric_mha_data_root,
)
StageII_Dataset = dict(
    type=mgam_SemiSup_3D_Mha,
    dataset_name="StageII",
    data_root=StageII_mha_data_root,
    data_root_mha=StageII_mha_data_root,
)
CPTAC_SAR_Dataset = dict(
    type=mgam_SemiSup_3D_Mha,
    dataset_name="CPTAC_SAR",
    data_root=CPTAC_SAR_mha_data_root,
    data_root_mha=CPTAC_SAR_mha_data_root,
)
Colonography_Dataset = dict(
    type=mgam_SemiSup_3D_Mha,
    dataset_name="Colonography",
    data_root=Colonography_mha_data_root,
    data_root_mha=Colonography_mha_data_root,
)
CT_ORG_Dataset = dict(
    type=mgam_SemiSup_3D_Mha,
    dataset_name="CT_ORG",
    data_root=CT_ORG_mha_data_root,
    data_root_mha=CT_ORG_mha_data_root,
)

TCGA_Datasets = [
    dict(
        type=mgam_SemiSup_3D_Mha,
        dataset_name=dataset_name,
        data_root=dataset_path,
        data_root_mha=dataset_path,
    )
    for dataset_name, dataset_path in TCGA_mha_data_roots.items()
]
